package com.jian.socks5.udp.handler;

import com.jian.socks5.common.Constants;
import com.jian.socks5.udp.DefaultSocks5UDPPacket;
import com.jian.socks5.udp.UDPBootstrap;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.channel.*;
import io.netty.channel.socket.DatagramPacket;

import java.net.InetSocketAddress;

@ChannelHandler.Sharable
public class Socks5UDPHandler extends SimpleChannelInboundHandler<DefaultSocks5UDPPacket> {

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, DefaultSocks5UDPPacket msg) throws Exception {
        String dstAddr = msg.getDstAddr();
        int dstPort = msg.getDstPort();
        ByteBuf data = msg.getData();

        InetSocketAddress sender = msg.getSender();

        //如果发送者端口为0，则不需要接受服务端的udp数据，只负责发送
        Bootstrap bootstrap;

        //转发的udp数据包
        DatagramPacket dirDatagramPacket = new DatagramPacket(data, new InetSocketAddress(dstAddr, dstPort));

        //客户端数据来源IP和端口和目标地址端口key
        String channelKey = sender.getHostName() + ":" + sender.getPort() + "_" + dstAddr + ":" + dstPort;

        Channel channel1 = Constants.dstIPAndPortEqSenderIPPort.get(channelKey);

        //channel复用，如果不存在，则需要重新连接，存在则直接使用该channel发送数据
        if(channel1 == null){
            //需要接受来自服务端的数据
            bootstrap = UDPBootstrap.newUDPInstance();
            ChannelFuture bind = bootstrap.bind(0);
            bind.addListener(future -> {
                ChannelFuture channelFuture = (ChannelFuture) future;
                if (channelFuture.isSuccess()) {
                    //设置需要接受数据的客户端ip信息
                    Channel channel = channelFuture.channel();
                    //绑定UDP后添加
                    Constants.dstIPAndPortEqSenderIPPort.put(channelKey, channel);
                    //
                    if(sender != null && sender.getPort() != 0){
                        channel.attr(Constants.udpClientSenderAddrAttr).set(sender);
                        channel.attr(Constants.udpClientSenderAddrTypeAttr).set(msg.getSocks5AddressType());
                        channel.attr(Constants.udpClientSenderChannelAttr).set(ctx);
                    }
                    //设置本机代理服务需要接受服务数据的端口，端口就为绑定的端口号
                    ChannelFuture channelFuture1 = channel.writeAndFlush(dirDatagramPacket);
                    channelFuture1.addListener(future1 -> {
                        if (future1.isSuccess()) {
                            System.out.println("UDP客户端数据转发至服务端成功！");
                        }else{
                            System.err.println("UDP客户端数据转发至服务端失败！");
                        }
                    });
                }else{
                    System.err.println("启动发送至目标机器UDP失败！");
                }
            });
        }else{
            //如果channel已经存在，则直接使用
            ChannelFuture channelFuture = channel1.writeAndFlush(dirDatagramPacket);
            channelFuture.addListener(future1 -> {
                if (future1.isSuccess()) {
                    System.out.println("UDP客户端数据转发至服务端成功！");
                }else{
                    System.err.println("UDP客户端数据转发至服务端失败！");
                }
            });
        }

    }
}
